# -*- coding: utf-8 -*-
# Author   : ZhangQing
# Time     : 2025-08-04 23:34
# File     : code_generator.py
# Project  : codebuddy_craft
# Desc     : 代码生成智能体
# backend/core/agents/code_generator.py
from typing import Dict, List, Optional, Any
from dataclasses import dataclass
import asyncio

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field

from backend.core.knowledge.code_templates import CodeTemplateManager
from backend.core.knowledge.language_specs import LanguageSpecManager
from backend.utils.config import get_settings


class CodeGenerationRequest(BaseModel):
    """代码生成请求模型"""
    prompt: str = Field(description="自然语言需求描述")
    tech_stack: str = Field(description="技术栈", example="FastAPI")
    context: Optional[str] = Field(default=None, description="上下文代码")
    model: Optional[str] = Field(default="gpt-4", description="使用的模型")


class CodeGenerationResponse(BaseModel):
    """代码生成响应模型"""
    code_block: str = Field(description="生成的代码")
    explanation: str = Field(description="代码说明")
    model_used: str = Field(description="使用的模型")
    suggestions: List[str] = Field(default=[], description="优化建议")


@dataclass
class AgentState:
    """智能体状态"""
    messages: List[BaseMessage]
    request: CodeGenerationRequest
    response: Optional[CodeGenerationResponse] = None
    context_analysis: Optional[Dict[str, Any]] = None
    template_match: Optional[Dict[str, Any]] = None
    error: Optional[str] = None


class CodeGeneratorAgent:
    """代码生成智能体"""

    def __init__(self):
        self.settings = get_settings()
        # 提供默认API密钥，避免验证错误
        api_key = self.settings.OPENAI_API_KEY or "sk-dummy-api-key-for-testing"
        self.llm = ChatOpenAI(
            model=self.settings.OPENAI_MODEL,
            temperature=0.1,
            api_key=api_key
        )
        self.template_manager = CodeTemplateManager()
        self.language_manager = LanguageSpecManager()
        self.graph = self._build_graph()

    def _build_graph(self) -> StateGraph:
        """构建智能体工作流图"""
        workflow = StateGraph(AgentState)

        # 添加节点
        workflow.add_node("analyze_context", self._analyze_context)
        workflow.add_node("match_template", self._match_template)
        workflow.add_node("generate_code", self._generate_code)
        workflow.add_node("optimize_code", self._optimize_code)
        workflow.add_node("format_response", self._format_response)

        # 设置入口点
        workflow.set_entry_point("analyze_context")

        # 添加边
        workflow.add_edge("analyze_context", "match_template")
        workflow.add_edge("match_template", "generate_code")
        workflow.add_edge("generate_code", "optimize_code")
        workflow.add_edge("optimize_code", "format_response")
        workflow.add_edge("format_response", END)

        return workflow.compile()

    async def _analyze_context(self, state: AgentState) -> AgentState:
        """分析上下文信息"""
        try:
            # 分析技术栈和需求
            analysis_prompt = ChatPromptTemplate.from_template("""
            分析以下编程需求，提取关键信息：

            需求描述：{prompt}
            技术栈：{tech_stack}
            上下文代码：{context}

            请分析：
            1. 功能目标是什么？
            2. 需要用到哪些技术组件？
            3. 有哪些约束条件？
            4. 输入输出要求是什么？

            以JSON格式返回分析结果。
            """)

            chain = analysis_prompt | self.llm
            response = await chain.ainvoke({
                "prompt": state.request.prompt,
                "tech_stack": state.request.tech_stack,
                "context": state.request.context or "无"
            })

            # 解析分析结果
            state.context_analysis = {
                "goal": "解析中...",
                "components": [],
                "constraints": [],
                "io_requirements": "分析中..."
            }

            return state

        except Exception as e:
            state.error = f"上下文分析失败: {str(e)}"
            return state

    async def _match_template(self, state: AgentState) -> AgentState:
        """匹配代码模板"""
        try:
            # 根据技术栈和需求匹配最佳模板
            template = await self.template_manager.find_best_template(
                tech_stack=state.request.tech_stack,
                description=state.request.prompt
            )

            state.template_match = {
                "template_id": template.get("id", "custom"),
                "template_code": template.get("code", ""),
                "template_description": template.get("description", ""),
                "match_score": template.get("score", 0.0)
            }

            return state

        except Exception as e:
            state.error = f"模板匹配失败: {str(e)}"
            return state

    async def _generate_code(self, state: AgentState) -> AgentState:
        """生成代码"""
        try:
            # 构建代码生成提示
            code_prompt = ChatPromptTemplate.from_template("""
            作为一个高级{tech_stack}开发工程师，请根据以下需求生成高质量的代码：

            需求描述：{prompt}
            技术栈：{tech_stack}

            {template_section}

            {context_section}

            请生成：
            1. 完整可运行的代码
            2. 遵循最佳实践和代码规范
            3. 包含必要的注释
            4. 考虑错误处理

            返回格式：
            ```{language}
            // 生成的代码
            ```

            代码说明：[简要说明代码的功能和使用方法]
            """)

            # 准备模板部分
            template_section = ""
            if state.template_match and state.template_match["template_code"]:
                template_section = f"""
                参考模板：
                ```
                {state.template_match["template_code"]}
                ```
                """

            # 准备上下文部分
            context_section = ""
            if state.request.context:
                context_section = f"""
                现有代码上下文：
                ```
                {state.request.context}
                ```
                """

            # 获取语言规范
            language_spec = await self.language_manager.get_language_info(
                state.request.tech_stack
            )

            chain = code_prompt | self.llm
            response = await chain.ainvoke({
                "prompt": state.request.prompt,
                "tech_stack": state.request.tech_stack,
                "template_section": template_section,
                "context_section": context_section,
                "language": language_spec.get("extension", "python")
            })

            # 解析生成的代码和说明
            content = response.content
            code_block = self._extract_code_block(content)
            explanation = self._extract_explanation(content)

            state.response = CodeGenerationResponse(
                code_block=code_block,
                explanation=explanation,
                model_used=state.request.model,
                suggestions=[]
            )

            return state

        except Exception as e:
            state.error = f"代码生成失败: {str(e)}"
            return state

    async def _optimize_code(self, state: AgentState) -> AgentState:
        """优化代码"""
        try:
            if not state.response:
                return state

            # 生成优化建议
            optimization_prompt = ChatPromptTemplate.from_template("""
            请分析以下代码并提供优化建议：

            ```{language}
            {code}
            ```

            请从以下方面提供建议：
            1. 性能优化
            2. 代码可读性
            3. 安全性
            4. 最佳实践

            返回3-5条具体的优化建议。
            """)

            language_spec = await self.language_manager.get_language_info(
                state.request.tech_stack
            )

            chain = optimization_prompt | self.llm
            response = await chain.ainvoke({
                "code": state.response.code_block,
                "language": language_spec.get("extension", "python")
            })

            # 解析优化建议
            suggestions = self._parse_suggestions(response.content)
            state.response.suggestions = suggestions

            return state

        except Exception as e:
            state.error = f"代码优化失败: {str(e)}"
            return state

    async def _format_response(self, state: AgentState) -> AgentState:
        """格式化响应"""
        try:
            if state.response:
                # 确保代码格式正确
                formatted_code = await self._format_code(
                    state.response.code_block,
                    state.request.tech_stack
                )
                state.response.code_block = formatted_code

            return state

        except Exception as e:
            state.error = f"响应格式化失败: {str(e)}"
            return state

    def _extract_code_block(self, content: str) -> str:
        """提取代码块"""
        import re

        # 匹配 ```language 代码块
        code_pattern = r'```[\w]*\n(.*?)\n```'
        matches = re.findall(code_pattern, content, re.DOTALL)

        if matches:
            return matches[0].strip()

        # 如果没有代码块标记，返回原内容
        return content.strip()

    def _extract_explanation(self, content: str) -> str:
        """提取代码说明"""
        import re

        # 查找"代码说明"部分
        explanation_pattern = r'代码说明[：:](.*?)(?=\n\n|\n```|$)'
        matches = re.findall(explanation_pattern, content, re.DOTALL)

        if matches:
            return matches[0].strip()

        # 如果没有找到说明，返回默认说明
        return "代码已生成，请查看上方代码块。"

    def _parse_suggestions(self, content: str) -> List[str]:
        """解析优化建议"""
        import re

        # 查找编号列表
        suggestions = re.findall(r'\d+\.\s*(.*?)(?=\n\d+\.|\n\n|$)', content, re.DOTALL)

        if suggestions:
            return [s.strip() for s in suggestions]

        # 如果没有找到编号列表，按行分割
        lines = content.strip().split('\n')
        return [line.strip() for line in lines if line.strip()][:5]

    async def _format_code(self, code: str, tech_stack: str) -> str:
        """格式化代码"""
        try:
            # 根据技术栈选择格式化器
            if "python" in tech_stack.lower():
                # 使用black格式化Python代码
                import black
                return black.format_str(code, mode=black.FileMode())
            elif "javascript" in tech_stack.lower() or "typescript" in tech_stack.lower():
                # 对于JS/TS，暂时返回原代码
                return code
            else:
                return code
        except Exception:
            # 格式化失败时返回原代码
            return code

    async def generate_code(self, request: CodeGenerationRequest) -> CodeGenerationResponse:
        """生成代码的主入口"""
        try:
            # 创建初始状态
            initial_state = AgentState(
                messages=[HumanMessage(content=request.prompt)],
                request=request
            )

            # 运行工作流
            final_state = await self.graph.ainvoke(initial_state)

            # 检查是否有错误
            if final_state.error:
                raise Exception(final_state.error)

            return final_state.response

        except Exception as e:
            # 返回错误响应
            return CodeGenerationResponse(
                code_block=f"# 代码生成失败\n# 错误: {str(e)}",
                explanation=f"抱歉，代码生成过程中出现错误：{str(e)}",
                model_used=request.model,
                suggestions=["请检查输入参数", "尝试简化需求描述", "确认技术栈正确"]
            )


# 创建全局实例
code_generator = CodeGeneratorAgent()

